import torch
from torch.utils.tensorboard import SummaryWriter
import os
import shutil
from collections import namedtuple
import random
import numpy as np
import gym

env_name = ['CartPole-v0', # 倒立摆，离散动作空间
            'Pendulum-v0', # 旋转摆，连续动作空间
            ]
env0 = gym.make(env_name[0]).unwrapped # unwrapped是解除每轮的步数限制
env1 = gym.make(env_name[1]).unwrapped # unwrapped是解除每轮的步数限制

class Memory(object):
    def __init__(self,capacity=10000,logs='./logs'):
        self.memory = []
        self.capacity = capacity
        self.step = 0
        self.Transition = namedtuple('transition',['s','a','r','s_','done'])

        if os.path.exists(logs):
            shutil.rmtree(logs)
        self.writer = SummaryWriter(logs)

        self.device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

    def put_transition(self,*args):
        """存储一组数据，包含：['s','a','r','s_','done']"""
        t = self.Transition(*args)
        if self.step < self.capacity:
            self.memory.append(t)
            self.step += 1
        else:
            idx = self.step % self.capacity
            self.memory[idx] = t
            self.step += 1
        return

    def get_transition(self,batch_size=1):
        """获取一批样本数据"""
        batch_size = min(batch_size,self.step)
        batch = random.sample(self.memory,k=batch_size)
        batch = self.Transition(*zip(*batch))
        return batch

    def importance_sample(self,batch_size=1):
        return


